/// 带取模的乘. a * b % modulus.
fn mul_mod(a: u128, b: u128, modulus: u128) -> u128 {
    // let (a, b, modulus) = (U256::from(a), U256::from(b), U256::from(modulus));
    // (a * b % modulus).as_u128()
    a * b % modulus
}

/// 带取模的快速幂. base^exponent % modulus.
///
/// 把exponent以二进制形式拆开. 结果为$\prod base^{2^exponent_i}$.
fn quick_pow_mod(mut base: u128, mut exponent: u128, modulus: u128) -> u128 {
    // let (base, exponent, modulus) = (U256::from(base), U256::from(exponent), U256::from(modulus));
    let mut ans = 1;
    while exponent > 0 {
        if exponent % 2 > 0 {
            ans = mul_mod(ans, base, modulus);
        }
        base = mul_mod(base, base, modulus);
        exponent /= 2;
    }
    ans
}

/// Miller–Rabin 素性测试.
pub fn miller_rabin(n: u128) -> bool {
    if n < 3 || n % 2 == 0 {
        return n == 2;
    }
    if n % 3 == 0 {
        return n == 3;
    }

    let mut u = n - 1;
    let mut t = 0;
    while u % 2 == 0 {
        u /= 2;
        t += 1;
    }

    const MR_TEST_TIME: u32 = 10;
    for _ in 0..MR_TEST_TIME {
        // a in [2, n-2]
        let a = fastrand::u128(2..=(n - 2));
        let mut v = quick_pow_mod(a, u, n);
        if v == 1 {
            continue;
        }

        let mut pass = false;
        for _ in 0..t {
            if v == n - 1 {
                pass = true;
                break;
            }
            v = mul_mod(v, v, n);
        }
        if !pass {
            return false;
        }
    }
    return true;
}
